package net

import (
	"context"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/google/gopacket/pcapgo"
	"io"
	"time"
)

type PcapHelp struct {
}

func NewPcapHelp() *PcapHelp {
	return &PcapHelp{}
}

func (tis *PcapHelp) Version() string {
	return pcap.Version()
}

func (tis *PcapHelp) EnumDevices() ([]pcap.Interface, error) {
	// 得到所有的(网络)设备
	devices, err := pcap.FindAllDevs()
	if err != nil {
		return nil, err
	}

	return devices, nil
}

func (tis *PcapHelp) Capture(ctx context.Context,
	deviceName string,
	onPacket func(packet gopacket.Packet),
	writer io.Writer,
	filter string) error {
	var (
		err         error
		snapshotLen uint32 = 100 // 20~60
		promiscuous        = true
		timeout            = -1 * time.Second
		handle      *pcap.Handle
		w           *pcapgo.Writer
	)

	if writer != nil {
		w = pcapgo.NewWriter(writer)
		_ = w.WriteFileHeader(snapshotLen, layers.LinkTypeEthernet)
	}

	// device: 网络设备的名称，如 eth0,也可以填充 pcap.FindAllDevs() 返回的设备的 Name
	// snaplen: 每个数据包读取的最大长度 the maximum size to read for each packet
	// promisc: 是否将网口设置为混杂模式, 即是否接收目的地址不为本机的包
	// timeout: 设置抓到包返回的超时。如果设置成 30s，那么每 30s 才会刷新一次数据包；设置成负数，会立刻刷新数据包，即不做等待
	handle, err = pcap.OpenLive(deviceName, int32(snapshotLen), promiscuous, timeout)
	if err != nil {
		return err
	}

	defer handle.Close()

	if len(filter) > 0 {
		// filter: "tcp || udp"
		if err = handle.SetBPFFilter(filter); err != nil {
			return err
		}
	}

	// Use the handle as a packet source to process all packets
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	for packet := range packetSource.Packets() {
		select {
		case <-ctx.Done():
			return nil
		default:
		}

		if onPacket != nil {
			onPacket(packet)
		}
		if w != nil {
			_ = w.WritePacket(packet.Metadata().CaptureInfo, packet.Data())
		}
	}

	return nil
}
